
import random
import math

'''
    test maximum entropy principle
'''


def compute_average(beta, amino_presence, energy_levels,  verbose=False):
    '''
        this function computes the expected energy
        given the parameter beta
    '''
    average=0.
    for amino, n_a in amino_presence.iteritems():
        expected_per_amino=0.
        normalization=0.

        for e, g_E in energy_levels[amino].iteritems():
            expected_per_amino += e * g_E * math.exp(-beta * e)
            normalization += g_E * math.exp(-beta * e)
        expected_per_amino/=normalization

        if verbose:
            print '=== amino', amino, 'Distribution (e, g_E, probability) ==='
            for e, g_E in energy_levels[amino].iteritems():
                print e, g_E, g_E * math.exp(-beta * e)/ normalization
            
        average += n_a * expected_per_amino

    return average


def beta_given_mu(mu, amino_presence, energy_levels):

    '''
        mu is the total GC content
        we infer beta so that mu matches 
        average(beta)
    '''
    rl = -10
    rr = 10
    avl=compute_average(rl, amino_presence, energy_levels)-mu
    avr=compute_average(rr, amino_presence, energy_levels)-mu
    
    #print 'initial limits:', avl/mu, avr/mu
    assert avl*avr<0, 'something wrong with '+str(mu)
    rm=0.5*(rl+rr)
    avv=compute_average(rm, amino_presence, energy_levels)-mu

    while math.fabs(avv)>1e-7:
        if avv*avl>0:
            rl=rm
        else:
            rr=rm
        rm=0.5*(rl+rr)
        avv=compute_average(rm, amino_presence, energy_levels)-mu
        #print rm, avv
    return rm




def get_a_sample():
    
    '''
        samples the genome one 
        uniformly at random
        (according to degeneracies)
        '''
    
    # samples[amino]=[sample]
    sample={}
    for amino, n_a in amino_presence.iteritems():
        #print 'amino', amino, 'degeneracies', energy_levels[amino]
        list_all=[]
        for e, g_E in energy_levels[amino].iteritems():
            for i in range(g_E):
                list_all.append(e)
        #print list_all
        am_sample=[]
        for i in range(n_a):
            am_sample.append(random.sample(list_all, 1)[0])
        sample[amino]=am_sample
    return sample



if __name__=='__main__':
    

    # IMPORTANT!!!
    # the analytical calculation
    # is correct only if we have 
    # enough nucleotides for each amino
    # (long genomes)
    # typically ~10 nucleotides per genome are
    # very good already
    
    # how much GC content is there in total
    # i.e. bases which are either G or C
    mu=21

    # for each amino-acid, this gives the 
    # degeneacies for several energy levels
    # energy_levels[amino] = { E: g_E }
    energy_levels={}

    # amino_presence[amino]=[num]
    amino_presence={}

    # here I define a synthetic case 
    amino_presence[1]=10
    amino_presence[2]=10
    amino_presence[3]=10
    energy_levels[1]={0:3, 1:2, 2:1}
    energy_levels[2]={0:1, 1:10}
    energy_levels[3]={0:1, 1:2, 3:1}
        
    N=sum(amino_presence.values())
        
    # analytical calculation
    beta = beta_given_mu(mu, amino_presence, energy_levels)
    print 'beta', beta
    
    random_e=[]
    samples=[]
    for i in range(10000):
        x=get_a_sample()
        sumx= 0
        for amino, sample in x.iteritems():
            sumx+=sum(sample)
        random_e.append(sumx)
        if sumx==mu:
            samples.append(x)
            #print len(samples), '<< new!'
    
    #print samples
    random_e_hist={}
    for e in random_e:
        random_e_hist[e]=random_e_hist.get(e,0)+1
    print '******', random_e_hist, '******'
    print 'random: ', float(sum(random_e))/len(random_e)

    # distr[amino] = histogram
    distr={}
    for amino in amino_presence:
        distr[amino]={}
    
    print 'good samples', len(samples)
    for sample in samples:
        for amino, s in sample.iteritems():
            for v in s:
                distr[amino][v]=distr[amino].get(v,0)+1

    for amino, d in distr.iteritems():
        all_counts=sum(d.values())
        print 'amino', amino, 'all_counts', all_counts
        for gc, v in d.iteritems():
            print gc, float(v)/all_counts

    compute_average(beta, amino_presence, energy_levels, verbose=True)
